package org.renwei.ml.bayes;

import com.google.common.collect.ComparisonChain;
import com.google.common.collect.Lists;
import com.google.common.collect.Multiset;
import org.renwei.ml.DataHelper;

import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * Created by renwei on 16/9/2.
 */
public class Bayes {
    private String decisionColumn;  //决定字段

    public Bayes(String decisionColumn) {
        this.decisionColumn = decisionColumn;
    }

    /**
     * @param dataSet 训练数据集
     * @param testMap 待测试数据集
     * @return
     */
    public String classify(Map<String, List<String>> dataSet, Map<String, String> testMap) {
        List<ColumnPro> columnPros = Lists.newArrayList();

        List<String> results = dataSet.get(decisionColumn);
        Multiset<String> multiset = DataHelper.getTypeCounts(results);
        for (Multiset.Entry<String> entry : multiset.entrySet()) {

            List<Integer> descIndex = DataHelper.getIndex(entry.getElement(), results);
            double pro = (double) entry.getCount() / dataSet.get(decisionColumn).size();
            for (String column : dataSet.keySet()) {
                if (column.equals(decisionColumn)) {
                    continue;
                }
                List<String> value = DataHelper.getSubSet(descIndex, dataSet.get(column));
                double testValueCount = null == testMap.get(column) ?
                    1.f :
                    DataHelper.getIndex(testMap.get(column), value).size();
                pro *= testValueCount / (1 + value.size());
            }

            columnPros.add(new ColumnPro(entry.getElement(), pro));
        }

        Collections.sort(columnPros);
        return columnPros.get(0).getColumn();
    }
}


class ColumnPro implements Comparable<ColumnPro> {
    private String column;
    private double pro;

    public ColumnPro(String column, double pro) {
        this.column = column;
        this.pro = pro;
    }

    public String getColumn() {
        return column;
    }

    public void setColumn(String column) {
        this.column = column;
    }

    public double getPro() {
        return pro;
    }

    public void setPro(double pro) {
        this.pro = pro;
    }

    public int compareTo(ColumnPro o) {
        return ComparisonChain.start().compare(o.pro, this.pro).result();
    }
}
